#include "mlir/Dialect/Func/Extensions/AllExtensions.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Dialect/Affine/Passes.h"
#include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
#include "mlir/ExecutionEngine/ExecutionEngine.h"
#include "mlir/ExecutionEngine/OptUtils.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Verifier.h"
#include "mlir/InitAllDialects.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Export.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"

#include "llvm/ADT/StringRef.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/ErrorOr.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"

#include "myDialect/toyDialect.h"
#include "myDialect/toyOps.h"

#include "toyLexer.h"
#include "toyParser.h"
#include "parser/MLIRToyVisitor.h"
#include "rewrite/passes.h"

#include <cassert>
#include <memory>
#include <string>
#include <system_error>
#include <utility>



int main(int argc, char **argv){
    std::string inputFileName = "/root/mlir_test/lib/parser/test.toy";
    std::fstream in(inputFileName);
    antlr4::ANTLRInputStream input(in);
    toyLexer lexer(&input);
    antlr4::CommonTokenStream tokens(&lexer);
    toyParser parser(&tokens);
    auto root = parser.module();

    mlir::MLIRContext context;
    context.getOrLoadDialect<mlir::toy::ToyDialect>();
    
    MLIRToyVisitor visitor(inputFileName, context);
    visitor.visit(root);
    
    mlir::ModuleOp mlirModule = visitor.getModule();
    mlirModule->dump();

    mlir::PassManager pm(&context);
    mlir::applyPassManagerCLOptions(pm);
    pm.addPass(mlir::toy::createLowerToArithPass());
    pm.addPass(mlir::toy::createLowerToLLVMPass());
    
    if (mlir::failed(pm.run(mlirModule)))
        return 4;
    std::cout << "======rewrited mlir======" << std::endl;
    mlirModule->dump();

    mlir::registerBuiltinDialectTranslation(context);
    mlir::registerLLVMDialectTranslation(context);

    llvm::LLVMContext llvmContext;
    auto llvmModule = mlir::translateModuleToLLVMIR(mlirModule, llvmContext);
    if(!llvmModule){
        llvm::errs() << "Failed to emit LLVM IR\n";
        return -1;
    }
    std::cout << "======llvm ir======" << std::endl;
    llvmModule->dump();
    return 0;
}